#type vertex
#version 330 core
layout (location = 0) in vec3 aPos;
layout (location = 1) in vec3 aNor;
layout (location = 2) in vec2 aTex;

out vec3 FragPos;
out vec3 FragNor;
out vec2 FragTex;
out vec4 pos_lightSpace;

uniform mat4 _model;
uniform mat4 _view;
uniform mat4 _projection;

uniform mat4 _light_space_matrix;

void main()
{
    pos_lightSpace = _light_space_matrix * _model * vec4(aPos, 1.0);
    FragPos = (_model * vec4(aPos, 1.0)).xyz;
    FragNor = mat3(transpose(inverse(_model))) * aNor;
    FragTex = aTex;
    gl_Position = _projection * _view * _model * vec4(aPos, 1.0);
}

#type fragment
#version 330 core

in vec3 FragPos;
in vec3 FragNor;
in vec2 FragTex;
in vec4 pos_lightSpace;

out vec4 FragColor;

////////////////
// Material
////////////////

uniform vec3 _eye_pos;

struct Material
{
    sampler2D diffuse;
    sampler2D specular;
    sampler2D emission;
    int shininess;
};

uniform Material _material;

////////////////
// Light
////////////////

#define DIRECTION_LIGHT_MAX 3
struct DirectionLight
{
    float intensity;

    vec3 ambient;
    vec3 diffuse;
    vec3 specular;

    vec3 direction;
};
uniform int _direct_light_nums;
uniform DirectionLight _direct_lights[DIRECTION_LIGHT_MAX];

#define POINT_LIGHT_MAX 5
struct PointLight
{
    float intensity;

    vec3 ambient;
    vec3 diffuse;
    vec3 specular;

    vec3 position;
    float constant;
    float linear;
    float quadratic;
};
uniform int _point_light_nums;
uniform PointLight _point_lights[POINT_LIGHT_MAX];

#define SPOT_LIGHT_MAX 5
struct SpotLight
{
    float intensity;

    vec3 ambient;
    vec3 diffuse;
    vec3 specular;

    vec3 position;
    float constant;
    float linear;
    float quadratic;
    vec3 direction;//  direction of the spoting cone
    float innerCutOff;//  inner cut off
    float outerCutOff;//  outer cut off
};
uniform int _spot_light_nums;
uniform SpotLight _spot_lights[SPOT_LIGHT_MAX];

// shadow map
uniform sampler2D _shadow_map;
uniform vec2 _shadow_map_size;

vec3 phong_diffuse(vec3 diffuse_para, vec3 normal_para, vec3 light_dir_para)
{
    float diff = max(dot(normal_para, light_dir_para), 0.0);
    return diffuse_para * diff * vec3(texture(_material.diffuse, FragTex));
}

vec3 phong_specular(vec3 specular_para, vec3 normal_para, vec3 light_dir_para, vec3 eye_dir_para)
{
    vec3 reflectDir = reflect(-light_dir_para, normal_para);
    float spec = pow(max(dot(eye_dir_para, reflectDir), 0.0), _material.shininess);
    return specular_para * spec * vec3(texture(_material.specular, FragTex));
}

float unpack(vec4 rgbaDepth) {
    const vec4 bitShift = vec4(1.0, 1.0/256.0, 1.0/(256.0*256.0), 1.0/(256.0*256.0*256.0));
    return dot(rgbaDepth, bitShift);
}

#define PI 3.141592653589793
#define PI2 6.283185307179586
highp float rand_1to1(highp float x) {
    // -1 -1
    return fract(sin(x)*10000.0);
}

highp float rand_2to1(vec2 uv) {
    // 0 - 1
    const highp float a = 12.9898, b = 78.233, c = 43758.5453;
    highp float dt = dot(uv.xy, vec2(a, b)), sn = mod(dt, PI);
    return fract(sin(sn) * c);
}

#define DISK_SAMPLE_NUM 20
vec2 disk_samples[DISK_SAMPLE_NUM];

void uniform_disk_sample(vec2 seed){
    float rd_num = rand_2to1(seed);
    float sampleX = rand_1to1(rd_num);
    float sampleY = rand_1to1(sampleX);
    float angle = sampleX * PI2;
    float radius = sqrt(sampleY);
    for (int i=0;i<DISK_SAMPLE_NUM;i++){
        disk_samples[i]=vec2(cos(angle)*radius, sin(angle)*radius);
        sampleX = rand_1to1(sampleX);
        sampleY = rand_1to1(sampleX);
        angle = sampleX * PI2;
        radius = sqrt(sampleY);
    }
}

float bias = 0.01;

float visibility_shadowMap(vec4 projPosition)
{
    vec3 shadowCoord = projPosition.xyz / projPosition.w;
    shadowCoord = shadowCoord * 0.5 + 0.5;
    return shadowCoord.z < texture(_shadow_map, shadowCoord.xy).r + bias ? 1.0 : 0.0;
}

float visibility_PCF(vec4 projPosition)
{
    vec3 shadowCoord = projPosition.xyz / projPosition.w;
    shadowCoord = shadowCoord * 0.5 + 0.5;
    float mapDepth = 0.0;
    vec2 unit = 1.0 / _shadow_map_size;
    uniform_disk_sample(shadowCoord.xy);
    for (int k=0;k<DISK_SAMPLE_NUM;k++){
        mapDepth += shadowCoord.z < texture(_shadow_map, shadowCoord.xy +
            unit * disk_samples[k]).r + bias ? 1.0 : 0.0;
    }
    return mapDepth / float(DISK_SAMPLE_NUM);
}

uniform float _light_area;

float visibility_PCSS(vec4 projPosition)
{
    vec3 shadowCoord = projPosition.xyz / projPosition.w;
    shadowCoord = shadowCoord * 0.5 + 0.5;
    float cur_depth = shadowCoord.z;
    vec2 unit = 1.0 / _shadow_map_size;
    // blocker searching
    float avg_depth = 0.0;
    int block_count = 0;
    uniform_disk_sample(shadowCoord.xy);
    for (int k=0;k<DISK_SAMPLE_NUM;k++){
        float block_depth = texture(_shadow_map, shadowCoord.xy + 5.0*unit * disk_samples[k]).r + bias;
        if (cur_depth > block_depth){
            avg_depth += block_depth;
            block_count++;
        }
    }
    if (block_count == 0){
        avg_depth = -avg_depth;
    } else {
        avg_depth /= float(block_count);
    }
    // penumbra size
    float filter_area = 0.0;
    filter_area = (cur_depth - avg_depth) * _light_area / avg_depth;
    // filtering
    float mapDepth = 0.0;
    for (int k=0;k<DISK_SAMPLE_NUM;k++){
        mapDepth += cur_depth < texture(_shadow_map, shadowCoord.xy +
        filter_area * unit * disk_samples[k]).r + bias ? 1.0 : 0.0;
    }
    return mapDepth / float(DISK_SAMPLE_NUM);
}

void main()
{
    vec3 light_dir;
    vec3 normal = normalize(FragNor);
    vec3 eye_dir = normalize(_eye_pos - FragPos);
    vec3 result = vec3(0.0);
    vec3 ambient, diffuse, specular;

    float visiable;

    for (int i = 0; i < _direct_light_nums && i < DIRECTION_LIGHT_MAX; i++)
    {
        light_dir = normalize(_direct_lights[i].direction);
        bias = max(0.02 * (1.0 - dot(normal, light_dir)), 0.003);
        visiable = visibility_PCSS(pos_lightSpace);
        ambient = _direct_lights[i].ambient * vec3(texture(_material.diffuse, FragTex));
        diffuse = phong_diffuse(_direct_lights[i].diffuse, normal, light_dir);
        specular = phong_specular(_direct_lights[i].specular, normal, light_dir, eye_dir);
        result += ambient + visiable * _direct_lights[i].intensity * (diffuse + specular);
    }

    for (int i = 0; i < _point_light_nums && i < POINT_LIGHT_MAX; i++)
    {
        vec3 light_distance = _point_lights[i].position - FragPos;
        float distance = length(light_distance);
        light_dir = normalize(light_distance);
        float attenuation = 1.0 / (_point_lights[i].constant + _point_lights[i].linear * distance
        + _point_lights[i].quadratic * (distance * distance));
        ambient = _point_lights[i].ambient * vec3(texture(_material.diffuse, FragTex));
        diffuse = phong_diffuse(_point_lights[i].diffuse, normal, light_dir);
        specular = phong_specular(_point_lights[i].specular, normal, light_dir, eye_dir);
        result += attenuation * (ambient + visiable * _point_lights[i].intensity * (diffuse + specular));
    }

    for (int i = 0; i < _spot_light_nums && i < SPOT_LIGHT_MAX; i++)
    {
        vec3 light_distance = _spot_lights[i].position - FragPos;
        float distance = length(light_distance);
        light_dir = normalize(light_distance);
        float attenuation = 1.0 / (_spot_lights[i].constant + _spot_lights[i].linear * distance
        + _spot_lights[i].quadratic * (distance * distance));
        ambient = _spot_lights[i].ambient * vec3(texture(_material.diffuse, FragTex));
        diffuse = phong_diffuse(_spot_lights[i].diffuse, normal, light_dir);
        specular = phong_specular(_spot_lights[i].specular, normal, light_dir, eye_dir);
        float theta = dot(light_dir, normalize(_spot_lights[i].direction));
        float intensity =
        clamp((theta - _spot_lights[i].outerCutOff) /
        (_spot_lights[i].innerCutOff - _spot_lights[i].outerCutOff),
        0.0, 1.0);
        result += attenuation * (ambient + visiable * _spot_lights[i].intensity * intensity * (diffuse + specular));
    }

    //    result += vec3(texture(_material.emission, FragTex));//  emission

    FragColor = vec4(result, 1.0);
}